//
//  KImageSegmentator.m
//  Kia
//
//  Created by Andrey on 27/03/2009.
//  Copyright 2009 Karma Software. All rights reserved.
//

#import "KImageSegmentator.h"
#import "KImage.h"
#import "KSimpleUIntArray.h"
#import "KSimpleFloatArray.h"
#import "KSegmentationMap.h"
#import "KPixel.h"

typedef struct _KSegment 
{
	CGFloat meanValue;
	NSUInteger pixelCount;
	CGFloat centroid;
} KSegment, *KSegmentRef, *KSegmentArray;

typedef struct _K2DSegment 
{
	CIELabPixelRep meanValue;
	NSUInteger pixelCount;
	CIELabPixelRep centroid;
} K2DSegment, *K2DSegmentArray;

CGFloat GetRandomFloat()
{
	return (CGFloat)random() / (CGFloat)RAND_MAX;
}

CGFloat GetRandomSymmetricFloat()
{
	CGFloat x = GetRandomFloat();
	
	return 2 * x - 1;
}

CGFloat GetHueDifference(NSPoint p1, NSPoint p2)
{
	CGFloat hue1 = p2.y / p2.x;
	CGFloat hue2 = p1.y / p1.x;
	
	return fabs(hue2 - hue1);
}

#define LIGHTNESS_WEIGHT 0.0

CGFloat GetMetricValue(CIELabPixelRep p1, CIELabPixelRep p2)
{
	// Convert to canonical a*b* rep
	p1.a -= 0.5; p1.b -= 0.5;
	p2.a -= 0.5; p2.b -= 0.5;
	
	CGFloat deltaL = p2.lightness - p1.lightness;
	CGFloat deltaA = p2.a - p1.a;
	CGFloat deltaB = p2.b - p1.b;
	
	return sqrt(LIGHTNESS_WEIGHT * deltaL * deltaL + deltaA * deltaA + deltaB * deltaB);
}

NSUInteger GetMinElementIndexInFloatArray(CGFloat* array, NSUInteger elementCount)
{
	CGFloat minValue = 1e30;
	NSUInteger minValueIndex = 0;
	
	for (int i = 0; i < elementCount; i++)
	{
		if (array[i] < minValue)
		{
			minValue = array[i];
			minValueIndex = i;
		}
	}
	
	return minValueIndex + 1;
}

@implementation KImageSegmentator

+ (KSegmentationMap*) segmentImage: (KImage*)image 
	   usingKMeansWithSegmentCount: (NSUInteger)segmentCount 
				   andUsingChannel: (enum KImageChannel)channel
{
	// Randomly choose initial segment centroids
	KSegmentArray segments = (KSegmentArray)malloc(sizeof(KSegment) * segmentCount);
	bzero(segments, sizeof(KSegment) * segmentCount);
	for (int i = 0; i < segmentCount; i++)
	{
		segments[i].centroid = (CGFloat)random() / (CGFloat)RAND_MAX;
//		segments[i].centroid = (CGFloat)i / (segmentCount - 1);
	}
	
	// Distances in hue space from current pixel to each of centroids
	CGFloat* distancesToCentroids = (CGFloat*)malloc(sizeof(CGFloat) * segmentCount);

	// Initial segmentation map
	NSUInteger imagePixelCount = image.width * image.height;
	KSimpleUIntArray* segmentationMap = [KSimpleUIntArray arrayWithCapacity:imagePixelCount];
	
	// Segmentation process
	Boolean segmentationDidChange = YES;
	while (segmentationDidChange) 
	{
		// Restore segment states
		segmentationDidChange = NO;
		NSLog(@"---------- New segmentation pass ------------");
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].meanValue = 0.0;
			segments[i].pixelCount = 0;
			NSLog(@"Segment #%i centroid:\t %.4f", i, segments[i].centroid);
		}
		
		// Compute new segmentation
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			KPixel* currentPixel = [image pixelAtX:x y:y];
				
			CGFloat channelValueAtCurrentPixel = 
			[currentPixel redundantColorVector].data[channel];
			
			for (int i = 0; i < segmentCount; i++) 
			{
				distancesToCentroids[i] = 
				fabs(channelValueAtCurrentPixel - segments[i].centroid);
			}
			
			NSUInteger closestCentroid = 
			GetMinElementIndexInFloatArray(distancesToCentroids, segmentCount);
			
			if (closestCentroid != segmentationMap.data[index])
				segmentationDidChange = YES;
			
			segmentationMap.data[index] = closestCentroid;
			
			segments[closestCentroid - 1].pixelCount++;
			segments[closestCentroid - 1].meanValue += channelValueAtCurrentPixel;
		}
		
		// Compute new centroids
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].centroid = segments[i].meanValue / segments[i].pixelCount;
		}
	}
	
	KSimpleFloatArray* distancesBetweenCentroids = 
	[KSimpleFloatArray arrayWithCapacity:(segmentCount - 1)];
	
	for (int i = 1; i < segmentCount; i++)
	{
		distancesBetweenCentroids.data[i - 1] = 
		fabs(segments[i].centroid - segments[i - 1].centroid);
	}
	
	CGFloat variance = [distancesBetweenCentroids maxValue];
	
	KSegmentationMap* result = 
	[[KSegmentationMap alloc] 
	 initWithPixelCount:imagePixelCount andSegmentCount:segmentCount];
	
	for (NSUInteger i = 0; i < segmentCount; i++)
	{
		CGFloat centroid = segments[i].centroid;
		
		// Compute probability map
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			KPixel* currentPixel = [image pixelAtX:x y:y];
			
			CGFloat channelValueAtCurrentPixel = 
			[currentPixel redundantColorVector].data[channel];
				
				CGFloat distanceToCentroid = centroid - channelValueAtCurrentPixel;
				CGFloat distanceToCentroidSquared = 
				distanceToCentroid * distanceToCentroid;
				
				result.segments[i].data[index] = 
				exp(-distanceToCentroidSquared / (2 * variance * variance));
		}
											 
		
		result.centroids.data[i] = centroid;
	}
	
	// Clean up
	free(distancesToCentroids);
	free(segments);
	
	return [result autorelease];
}

+ (K2DSegmentationMap*) segmentImageByChromaticity: (KImage*)image 
					   usingKMeansWithSegmentCount: (NSUInteger)segmentCount
{
	NSUInteger imagePixelCount = image.width * image.height;
	
	// Form final segmentation map 
	K2DSegmentationMap* result = 
	[[K2DSegmentationMap alloc] 
	 initWithPixelCount:imagePixelCount andSegmentCount:segmentCount];
	
	K2DSegmentArray segments = (K2DSegmentArray)malloc(sizeof(K2DSegment) * segmentCount);
	
	for (int i = 0; i < segmentCount; i++)
	{
		segments[i].centroid.lightness = (CGFloat)i / segmentCount;
		segments[i].centroid.a = 0.5 * (1 + 0.25 * GetRandomSymmetricFloat());
		segments[i].centroid.b = 0.5 * (1 + 0.25 * GetRandomSymmetricFloat());
	}
	
	// Distances to centroids
	KSimpleFloatArray* distancesToCentroids = 
	[KSimpleFloatArray arrayWithCapacity:segmentCount];

	
	Boolean segmentationDidChange = YES;
	while (segmentationDidChange) 
	{
		// Restore segment states
//		NSLog(@"---------- New segmentation pass ------------");
		segmentationDidChange = NO;
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].meanValue = MakeCIELabPixelRep(0.0, 0.0, 0.0);
			segments[i].pixelCount = 0;
			/*
			NSLog(@"Segment #%i centroid:\t {%.4f, %.4f, %.4f}", i, 
				  segments[i].centroid.lightness, 
				  segments[i].centroid.a,
				  segments[i].centroid.b
			);
			 */
		}
		
		// Compute new segmentation
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			
			KPixel* currentPixel = [image pixelAtX:x y:y];
			
			for (int i = 0; i < distancesToCentroids.count; i++)
			{
				distancesToCentroids.data[i] = 
				GetMetricValue([currentPixel cieLabRep], 
							   segments[i].centroid);
			}
			
			NSUInteger closestCentroid = [distancesToCentroids minValueIndex];
			assert(closestCentroid != -1);
			if (closestCentroid != result.crispSegmentationMap.data[index])
				segmentationDidChange = YES;
			
			result.crispSegmentationMap.data[index] = closestCentroid;
			
			CIELabPixelRep previousMeanValue = segments[closestCentroid].meanValue;
			segments[closestCentroid].meanValue =
			MakeCIELabPixelRep(previousMeanValue.lightness + [currentPixel lightness], 
							   previousMeanValue.a + [currentPixel a], 
							   previousMeanValue.b + [currentPixel b]);											
			segments[closestCentroid].pixelCount++;
		}
		
		// Compute new centroids
		
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].centroid =
			MakeCIELabPixelRep(segments[i].meanValue.lightness / segments[i].pixelCount, 
							   segments[i].meanValue.a / segments[i].pixelCount,
							   segments[i].meanValue.b / segments[i].pixelCount
						);
		}
	}
	
	KSimpleFloatArray* distancesBetweenCentroids = 
	[KSimpleFloatArray arrayWithCapacity:(segmentCount - 1)];
	
	for (int i = 1; i < segmentCount; i++)
	{
		distancesBetweenCentroids.data[i - 1] = 
		GetMetricValue(segments[i].centroid, segments[i - 1].centroid);
	}
	
	CGFloat variance = [distancesBetweenCentroids maxValue] / 2;
	
	for (NSUInteger i = 0; i < segmentCount; i++)
	{
		CIELabPixelRep centroid = segments[i].centroid;
		
		// Compute probability map
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			
			KPixel* currentPixel = [image pixelAtX:x y:y];
						
			CGFloat distanceToCentroid = 
			GetMetricValue([currentPixel cieLabRep], centroid);
			
			CGFloat distanceToCentroidSquared = 
			distanceToCentroid * distanceToCentroid;

			result.segments[i].data[index] = 
			exp(-distanceToCentroidSquared / (2 * variance * variance));

		/*	
			result.segments[i].data[index] = distanceToCentroid;
		 */
		}
		
		result.centroidPoints[i] = centroid;
	}
	
	/* Normalize probability map */

	ForEachImagePixel(image)
	{
		NSUInteger index = y * image.width + x;
		CGFloat norm = 0.0;
		
		for (int i = 0; i < segmentCount; i++)
		{
			norm += result.segments[i].data[index];
		}
		
		for (int i = 0; i < segmentCount; i++)
		{
			result.segments[i].data[index] /= norm;
		}
	}

	free(segments);
	
	return [result autorelease];
}

@end
